#include "BMPSaver.h"
#include "Utils.h"
#include <iostream>
#include <shlwapi.h>

#pragma comment(lib, "shlwapi.lib")

BMPSaver::BMPSaver() {
}

BMPSaver::~BMPSaver() {
}

bool BMPSaver::CreateDirectoryIfNotExists(const std::string& dirPath) {
    if (dirPath.empty()) {
        m_lastError = "Empty directory path";
        return false;
    }

    if (PathFileExistsA(dirPath.c_str())) {
        DWORD attributes = GetFileAttributesA(dirPath.c_str());
        if (attributes != INVALID_FILE_ATTRIBUTES && (attributes & FILE_ATTRIBUTE_DIRECTORY)) {
            return true;
        } else {
            m_lastError = "Path exists but is not a directory: " + dirPath;
            return false;
        }
    }

    if (!CreateDirectoryA(dirPath.c_str(), NULL)) {
        DWORD error = ::GetLastError();
        if (error == ERROR_PATH_NOT_FOUND) {
            size_t lastSlash = dirPath.find_last_of("\\/");
            if (lastSlash != std::string::npos) {
                std::string parentDir = dirPath.substr(0, lastSlash);
                if (!CreateDirectoryIfNotExists(parentDir)) {
                    return false;
                }

                if (!CreateDirectoryA(dirPath.c_str(), NULL)) {
                    m_lastError = "Failed to create directory after creating parent: " + dirPath + " - " + Utils::GetLastWin32Error();
                    return false;
                }
            } else {
                m_lastError = "Failed to create directory: " + dirPath + " - " + Utils::GetLastWin32Error();
                return false;
            }
        } else {
            m_lastError = "Failed to create directory: " + dirPath + " - " + Utils::GetLastWin32Error();
            return false;
        }
    }

    return true;
}

bool BMPSaver::SaveBitmap(HBITMAP hBitmap, const std::string& filePath) {
    if (!hBitmap) {
        m_lastError = "Invalid bitmap handle";
        return false;
    }

    if (filePath.empty()) {
        m_lastError = "Empty file path";
        return false;
    }

    size_t lastSlash = filePath.find_last_of("\\/");
    if (lastSlash != std::string::npos) {
        std::string dirPath = filePath.substr(0, lastSlash);
        if (!CreateDirectoryIfNotExists(dirPath)) {
            return false;
        }
    }

    return WriteBitmapToFile(hBitmap, filePath);
}

bool BMPSaver::WriteBitmapToFile(HBITMAP hBitmap, const std::string& filePath) {
    BITMAP bitmap;
    if (!GetObject(hBitmap, sizeof(BITMAP), &bitmap)) {
        m_lastError = "Failed to get bitmap object: " + Utils::GetLastWin32Error();
        return false;
    }

    BITMAPINFO bitmapInfo;
    ZeroMemory(&bitmapInfo, sizeof(BITMAPINFO));
    bitmapInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bitmapInfo.bmiHeader.biWidth = bitmap.bmWidth;
    bitmapInfo.bmiHeader.biHeight = bitmap.bmHeight;
    bitmapInfo.bmiHeader.biPlanes = 1;
    bitmapInfo.bmiHeader.biBitCount = 24;
    bitmapInfo.bmiHeader.biCompression = BI_RGB;
    bitmapInfo.bmiHeader.biSizeImage = 0;

    HDC hdc = GetDC(NULL);
    if (!hdc) {
        m_lastError = "Failed to get device context: " + Utils::GetLastWin32Error();
        return false;
    }

    LPVOID pBits = NULL;
    DWORD bitmapSize = GetBitmapBits(hBitmap, &pBits, &bitmapInfo);
    if (bitmapSize == 0 || !pBits) {
        ReleaseDC(NULL, hdc);
        return false;
    }

    BITMAPFILEHEADER bitmapFileHeader;
    ZeroMemory(&bitmapFileHeader, sizeof(BITMAPFILEHEADER));
    bitmapFileHeader.bfType = 0x4D42;
    bitmapFileHeader.bfSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + bitmapSize;
    bitmapFileHeader.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);

    HANDLE hFile = CreateFileA(filePath.c_str(), GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        m_lastError = "Failed to create file: " + filePath + " - " + Utils::GetLastWin32Error();
        delete[] pBits;
        ReleaseDC(NULL, hdc);
        return false;
    }

    DWORD bytesWritten;
    bool success = true;

    if (!WriteFile(hFile, &bitmapFileHeader, sizeof(BITMAPFILEHEADER), &bytesWritten, NULL) ||
        bytesWritten != sizeof(BITMAPFILEHEADER)) {
        m_lastError = "Failed to write bitmap file header: " + Utils::GetLastWin32Error();
        success = false;
    }

    if (success && (!WriteFile(hFile, &bitmapInfo.bmiHeader, sizeof(BITMAPINFOHEADER), &bytesWritten, NULL) ||
        bytesWritten != sizeof(BITMAPINFOHEADER))) {
        m_lastError = "Failed to write bitmap info header: " + Utils::GetLastWin32Error();
        success = false;
    }

    if (success && (!WriteFile(hFile, pBits, bitmapSize, &bytesWritten, NULL) ||
        bytesWritten != bitmapSize)) {
        m_lastError = "Failed to write bitmap data: " + Utils::GetLastWin32Error();
        success = false;
    }

    CloseHandle(hFile);
    delete[] pBits;
    ReleaseDC(NULL, hdc);

    return success;
}

DWORD BMPSaver::GetBitmapBits(HBITMAP hBitmap, LPVOID* ppvBits, PBITMAPINFO pbi) {
    BITMAP bitmap;
    if (!GetObject(hBitmap, sizeof(BITMAP), &bitmap)) {
        m_lastError = "Failed to get bitmap object: " + Utils::GetLastWin32Error();
        return 0;
    }

    WORD bitsPerPixel = pbi->bmiHeader.biBitCount;
    DWORD bytesPerLine = ((bitmap.bmWidth * bitsPerPixel + 31) / 32) * 4;
    DWORD bitmapSize = bytesPerLine * bitmap.bmHeight;

    pbi->bmiHeader.biSizeImage = bitmapSize;

    *ppvBits = new BYTE[bitmapSize];
    if (!*ppvBits) {
        m_lastError = "Failed to allocate memory for bitmap data";
        return 0;
    }

    HDC hdc = GetDC(NULL);
    if (!hdc) {
        m_lastError = "Failed to get device context: " + Utils::GetLastWin32Error();
        delete[] *ppvBits;
        *ppvBits = NULL;
        return 0;
    }

    int result = GetDIBits(hdc, hBitmap, 0, bitmap.bmHeight, *ppvBits, pbi, DIB_RGB_COLORS);
    if (result == 0) {
        m_lastError = "Failed to get DIB bits: " + Utils::GetLastWin32Error();
        delete[] *ppvBits;
        *ppvBits = NULL;
        ReleaseDC(NULL, hdc);
        return 0;
    }

    ReleaseDC(NULL, hdc);
    return bitmapSize;
}